"""
Kafka相关数据模型
"""
from typing import Optional, Dict, Any, List
from pydantic import BaseModel, Field
from datetime import datetime

class KafkaConnectionConfig(BaseModel):
    """Kafka连接配置"""
    bootstrap_servers: str = Field(..., description="Kafka服务器地址")
    security_protocol: str = Field(default="PLAINTEXT", description="安全协议")
    sasl_mechanism: Optional[str] = Field(default=None, description="SASL认证机制")
    sasl_plain_username: Optional[str] = Field(default=None, description="SASL用户名")
    sasl_plain_password: Optional[str] = Field(default=None, description="SASL密码")
    
    class Config:
        schema_extra = {
            "example": {
                "bootstrap_servers": "localhost:9092",
                "security_protocol": "SASL_PLAINTEXT",
                "sasl_mechanism": "PLAIN",
                "sasl_plain_username": "admin",
                "sasl_plain_password": "password"
            }
        }

class TopicInfo(BaseModel):
    """Topic信息"""
    name: str = Field(..., description="Topic名称")
    partition_count: int = Field(..., description="分区数")
    replication_factor: int = Field(..., description="副本因子")
    configs: Dict[str, Any] = Field(default_factory=dict, description="Topic配置")
    
    class Config:
        schema_extra = {
            "example": {
                "name": "test-topic",
                "partition_count": 3,
                "replication_factor": 1,
                "configs": {
                    "retention.ms": "604800000",
                    "cleanup.policy": "delete"
                }
            }
        }

class PartitionInfo(BaseModel):
    """分区信息"""
    partition_id: int = Field(..., description="分区ID")
    leader: int = Field(..., description="Leader节点ID")
    replicas: List[int] = Field(..., description="副本节点列表")
    isr: List[int] = Field(..., description="同步副本列表")
    
    class Config:
        schema_extra = {
            "example": {
                "partition_id": 0,
                "leader": 1,
                "replicas": [1, 2],
                "isr": [1, 2]
            }
        }

class MessageRequest(BaseModel):
    """消息发送请求"""
    topic: str = Field(..., description="目标Topic")
    value: str = Field(..., description="消息内容")
    key: Optional[str] = Field(default=None, description="消息Key")
    partition: Optional[int] = Field(default=None, description="指定分区")
    headers: Optional[Dict[str, str]] = Field(default=None, description="消息头")
    
    class Config:
        schema_extra = {
            "example": {
                "topic": "test-topic",
                "value": '{"message": "Hello Kafka"}',
                "key": "user-123",
                "partition": 0,
                "headers": {"source": "web-ui"}
            }
        }

class BatchMessageRequest(BaseModel):
    """批量消息发送请求"""
    topic: str = Field(..., description="目标Topic")
    messages: List[Dict[str, Any]] = Field(..., description="消息列表")
    
    class Config:
        schema_extra = {
            "example": {
                "topic": "test-topic",
                "messages": [
                    {"value": "message1", "key": "key1"},
                    {"value": "message2", "key": "key2"}
                ]
            }
        }

class MessageResponse(BaseModel):
    """消息发送响应"""
    success: bool = Field(..., description="是否成功")
    partition: Optional[int] = Field(default=None, description="分区")
    offset: Optional[int] = Field(default=None, description="偏移量")
    timestamp: Optional[datetime] = Field(default=None, description="时间戳")
    error: Optional[str] = Field(default=None, description="错误信息")
    
    class Config:
        schema_extra = {
            "example": {
                "success": True,
                "partition": 0,
                "offset": 12345,
                "timestamp": "2024-01-01T10:00:00Z",
                "error": None
            }
        }

class ConnectionStatus(BaseModel):
    """连接状态"""
    connected: bool = Field(..., description="是否已连接")
    bootstrap_servers: Optional[str] = Field(default=None, description="连接的服务器地址")
    error: Optional[str] = Field(default=None, description="错误信息")
    connected_at: Optional[datetime] = Field(default=None, description="连接时间")
    
    class Config:
        schema_extra = {
            "example": {
                "connected": True,
                "bootstrap_servers": "localhost:9092",
                "error": None,
                "connected_at": "2024-01-01T10:00:00Z"
            }
        } 